source("/scratch/jc227089/evo-dispersal/evostoch/evostochFunctions.R")
setwd("/scratch/jc227089/evostochData/")

fname<-"evostoch"
flist<-list.files(pattern=fname)
flist<-flist[!grepl(pattern="concat", flist)]
flist<-flist[!grepl(pattern="Allee", flist)]

nreps<-20

concat<-c()
poplist<-vector(mode="list", length=nreps*length(flist))
for (ii in 1:length(flist)) {
	load(flist[ii])
	print(length(out))
	for (jj in 1:length(out)){
		if (is.null(out[[jj]]$pop)) next
		poplist[[(ii-1)*nreps+(jj)]]<-out[[jj]]$pop
		pars<-rep(unlist(out[[jj]]$parameters), each=2)
		pars<-matrix(pars, nrow=2)
		pars<-cbind(matrix(rep(c(pop=ii, rep=jj), each=2), nrow=2), pars, out[[jj]]$out)
		concat<-rbind(concat, pars)
	}
}
colnames(concat)<-c("pop", "rep", names(unlist(out[[jj]]$parameters)), 
		colnames(out[[jj]]$out))
colnames(concat)[16]<-'xlim'
concat<-cbind(concat, evol=concat[,"h2H"]>0, abs.xlim=abs(concat[,"xlim"]))


save(concat, poplist, file="concat_evostoch_varR0.RData")

c.df<-as.data.frame(concat)
rm(concat)

var.summ<-c()
for (ii in 1:nlevels(as.factor(c.df$R0))){
	c.tmp<-subset(c.df, as.numeric(as.factor(c.df$R0))==ii)
	c.no.evol<-subset(c.tmp, c.tmp$evol==0)
	c.evol<-subset(c.tmp, c.tmp$evol==1)
	
	dem.st<-summary(aov(abs.xlim~1, data=c.no.evol))[[1]]$'Sum Sq' #pure demographic stochasticity
	evol.st<-summary(aov(abs.xlim~Error(pop), data=c.evol)) # partitioned into within and between pop (between=founder, within = evol+demographic stoch)
	evol.st.win<-evol.st$'Error: Within'[[1]]$'Sum Sq' #evol+dem st
	evol.st.bw<-evol.st$'Error: pop'[[1]]$'Sum Sq' #initial founder event

	#proportion attributable to each effect
	dem.st.prop<-dem.st/(evol.st.win+evol.st.bw)
	found.st.prop<-evol.st.bw/(evol.st.win+evol.st.bw)
	evol.st.prop<-(evol.st.win-dem.st)/(evol.st.win+evol.st.bw)
	var.summ<-rbind(var.summ, c(R0=as.numeric(levels(as.factor(c.df$R0)))[ii], round(c(demog=dem.st.prop, founder=found.st.prop, evoln=evol.st.prop)*100, 2)))
}


####################################################
### Figures ###
c.evol<-subset(c.df, c.df$R0==16 & c.df$evol==1) # choose a single value of R0 for figures
nreps<-20

evens<-seq(2, nrow(c.evol), 2)
odds<-evens-1
xdiff<-c.evol$abs.xlim[evens]+c.evol$abs.xlim[odds]
sset<-sample(which(xdiff>0.95*max(xdiff)), 1)
popref<-c.evol[evens[sset] ,c('pop', 'rep')]

sample.pop<-poplist[[(popref$pop-1)*nreps+(popref$rep)]]
plotter.mean(sample.pop, a=c.evol$a[sset], R0=c.evol$R0[sset], filename='Figures/samplePop.pdf')

rel.fitness<-h.surv(1, c.evol$MeanH)
mod<-lm(c.evol$abs.xlim~c.evol$MeanD)
res.dist<-mod$residuals

pdf('Figures/correlations.pdf')
	par(mfrow=c(2,1), mar=c(5,5,1,1), cex.lab=1.4)
	plot(c.evol$MeanD, c.evol$abs.xlim, xlab="Evolved dispersal value on expanding front", ylab="Distance spread", pch=19, bty="l")
	
	plot(rel.fitness, res.dist, xlab="Relative fitness on expanding front", ylab="Residual distance spread", pch=19, bty="l")
dev.off()


var.summ[,2:4]<-var.summ[,2:4]/100
pdf('Figures/varR0.pdf')
	par(mar=c(5,5,1,1), cex.lab=1.4)
	matplot(x=var.summ[,1], y=var.summ[,2:4], type="l", lwd=3, col=1, bty="l",
		xlab=quote(italic(R[max])), ylab='Proportion of variance in spread distance')
	legend('right', legend=c('Demographic stochasticity', 'Initial founder event', 'Evolutionary stochasticity'), col=1, lty=1:3, lwd=3, bty='n' )
dev.off()